#! /usr/bin/env python3
from PIL import Image
from glob import glob
import os
from tqdm import tqdm
import argparse
import numpy as np

def resize(img, size):
    image = img
    fg_w, fg_h = image.size
    bg_w, bg_h = size
    
    if image.mode == 'L':
        out_image = np.zeros((bg_h, bg_w), dtype='uint8') + 255
    elif image.mode == 'RGB':
        out_image = np.zeros((bg_h, bg_w, 3), dtype='uint8') + 255
    elif image.mode == 'RGBA':
        out_image = np.zeros((bg_h, bg_w, 4), dtype='uint8') + 255
    else:
        image = image.convert('RGB')
        out_image = np.zeros((bg_h, bg_w, 3), dtype='uint8') + 255

    if fg_h == fg_w:
        image2 = image.resize(size, resample=Image.BOX)
        out_image = np.array(image2, dtype='uint8')
    elif fg_h > fg_w:
        h = bg_h
        w = int(h / fg_h * fg_w)
        image2 = image.resize((w, h), resample=Image.BOX)
        image2_arr = np.array(image2, dtype='uint8')
        out_image[:, (bg_w - w) // 2: (bg_w - w) // 2 + w] = image2_arr
    elif fg_w > fg_h:
        w = bg_w
        h = int(w / fg_w * fg_h)
        image2 = image.resize((w, h), resample=Image.BOX)
        image2_arr = np.array(image2, dtype='uint8')
        out_image[(bg_h - h) // 2: (bg_h - h) // 2 + h, :] = image2_arr
    return Image.fromarray(out_image)

def resizeImages(dir1, dir2, size):
    pathList = glob(os.path.join(dir1, '*'))
    for i in tqdm(range(len(pathList)), ncols=100):
        path1 = pathList[i]
        filename = os.path.split(path1)[-1]
        path2 = os.path.join(dir2, filename)
        image = resize(Image.open(path1), size)
        image.save(path2)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='图片（下）采样工具')
    parser.add_argument('source', metavar='原始图片目录', type=str,
                        help='原始图片文件所在目录')
    parser.add_argument('designation', metavar='生成图片目录', type=str,
                        help='生成图片文件所在目录')
    parser.add_argument('-s', dest='size', type=int, default=256,
                        help='图片缩放尺寸，默认 256×256')

    args = parser.parse_args()
    resizeImages(args.source, args.designation, (args.size, args.size))
